import numpy as np
import struct
import os

def deconvolution(input_data, weight, bias=None, stride=2, padding=1, output_padding=0):
    """
    Transposed convolution (deconvolution)
    
    Parameters:
    - input_data: (input_height, input_width, in_channels) - uint8
    - weight: (out_channels, in_channels, kernel_height, kernel_width) - uint8
    - stride: upsampling factor
    - padding: padding applied to input
    - output_padding: additional size added to output
    
    Output size: (IH - 1) * S + K - 2*P + OP
    """
    input_height, input_width, in_channels = input_data.shape
    out_channels, _, kernel_height, kernel_width = weight.shape
    
    # Calculate output dimensions
    out_height = (input_height - 1) * stride + kernel_height - 2 * padding + output_padding
    out_width = (input_width - 1) * stride + kernel_width - 2 * padding + output_padding
    
    output = np.zeros((out_height, out_width, out_channels), dtype=np.int32)
    
    for oc in range(out_channels):
        for oh in range(out_height):
            for ow in range(out_width):
                conv_sum = 0
                
                # For each output position, check all kernel positions
                for ic in range(in_channels):
                    for kh in range(kernel_height):
                        for kw in range(kernel_width):
                            # Calculate which input position contributes to this output
                            h_temp = oh - kh + padding
                            w_temp = ow - kw + padding
                            
                            # Check if this maps to a valid input position
                            if h_temp >= 0 and h_temp % stride == 0 and \
                               w_temp >= 0 and w_temp % stride == 0:
                                ih = h_temp // stride
                                iw = w_temp // stride
                                
                                if 0 <= ih < input_height and 0 <= iw < input_width:
                                    conv_sum += int(input_data[ih, iw, ic]) * int(weight[oc, ic, kh, kw])
                
                if bias is not None:
                    conv_sum += bias[oc]
                output[oh, ow, oc] = conv_sum
    
    return output


def reshape_weight_for_hardware(weight, P_OCH, P_ICH):
    """
    Reshape weight to match hardware format for deconv:
    Weight ROM: weight[FOLD_O][FOLD_I][K*K][P_OCH][P_ICH]
    
    Input weight shape: (N_OCH, N_ICH, K, K)
    Output shape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    
    Note: For deconv, kernels are flipped (reversed) in hardware
    """
    N_OCH, N_ICH, K, _ = weight.shape
    FOLD_O = N_OCH // P_OCH
    FOLD_I = N_ICH // P_ICH
    KK = K * K
    
    # Reshape: (N_OCH, N_ICH, K, K) -> (FOLD_O, P_OCH, FOLD_I, P_ICH, K, K)
    weight_reshaped = weight.reshape(FOLD_O, P_OCH, FOLD_I, P_ICH, K, K)
    
    # Transpose to: (FOLD_O, FOLD_I, K, K, P_OCH, P_ICH)
    weight_reshaped = weight_reshaped.transpose(0, 2, 4, 5, 1, 3)
    
    # Reshape to: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    weight_reshaped = weight_reshaped.reshape(FOLD_O, FOLD_I, KK, P_OCH, P_ICH)
    
    return weight_reshaped


def save_input_to_binary(input_data, filename, P_ICH, A_BIT=8):
    """
    Save input data to binary file for SystemVerilog testbench
    Input shape: (height, width, channels)
    Saved as packed vectors: each vector contains P_ICH channels
    
    For deconv, data is stored row by row, with channels folded into P_ICH groups
    """
    height, width, channels = input_data.shape
    
    with open(filename, 'wb') as f:
        for h in range(height):
            for w in range(width):
                for ch_start in range(0, channels, P_ICH):
                    # Pack P_ICH channels into one vector
                    vec = input_data[h, w, ch_start:ch_start+P_ICH]
                    for val in vec:
                        # Write as unsigned 8-bit integer
                        f.write(struct.pack('B', int(val) & 0xFF))
    
    print(f"Saved input (binary) to {filename}")
    print(f"  Shape: height={height}, width={width}, channels={channels}")
    print(f"  Vectors: {height * width * (channels // P_ICH)}")


def save_input_to_text(input_data, filename):
    """
    Save input data to text file
    """
    height, width, channels = input_data.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Deconv Input shape: height={height}, width={width}, channels={channels}\n")
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = input_data[h, w, c]
                    f.write(f"{int(val)}\n")
    
    print(f"Saved input (text) to {filename}")


def save_weight_to_binary(weight, filename, P_OCH, P_ICH, W_BIT=8):
    """
    Save weight data to binary file for SystemVerilog testbench
    Weight shape after reshape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    """
    FOLD_O, FOLD_I, KK, _, _ = weight.shape
    
    with open(filename, 'wb') as f:
        for fo in range(FOLD_O):
            for fi in range(FOLD_I):
                for kk in range(KK):
                    # Pack in order: [och0_ich0, och0_ich1, ..., och0_ichN, och1_ich0, ...]
                    for o in range(P_OCH):
                        for i in range(P_ICH):
                            val = weight[fo, fi, kk, o, i]
                            # Write as signed 8-bit integer
                            f.write(struct.pack('b', int(val)))
    
    print(f"Saved weight (binary) to {filename}")
    print(f"  Shape: FOLD_O={FOLD_O}, FOLD_I={FOLD_I}, K*K={KK}")
    print(f"  Total entries: {FOLD_O * FOLD_I * KK}")


def save_weight_to_text(weight, filename):
    """
    Save weight data to text file
    Weight shape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    """
    FOLD_O, FOLD_I, KK, P_OCH, P_ICH = weight.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Weight shape: FOLD_O={FOLD_O}, FOLD_I={FOLD_I}, K*K={KK}, P_OCH={P_OCH}, P_ICH={P_ICH}\n")
        for fo in range(FOLD_O):
            for fi in range(FOLD_I):
                for kk in range(KK):
                    for o in range(P_OCH):
                        for i in range(P_ICH):
                            val = weight[fo, fi, kk, o, i]
                            f.write(f"{int(val)}\n")
    
    print(f"Saved weight (text) to {filename}")


def save_output_to_binary(output, filename, P_OCH, B_BIT=32):
    """
    Save output data to binary file
    For deconv, output is saved in (height, width, channels) order
    Each P_OCH outputs are written together
    """
    height, width, channels = output.shape
    
    with open(filename, 'wb') as f:
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = output[h, w, c]
                    # Write as signed 32-bit integer (little-endian)
                    f.write(struct.pack('i', int(val)))
    
    print(f"Saved output (binary) to {filename}")
    print(f"  Shape: height={height}, width={width}, channels={channels}")


def save_output_to_text(output, filename):
    """
    Save output data to text file
    """
    height, width, channels = output.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Deconv Output shape: height={height}, width={width}, channels={channels}\n")
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = output[h, w, c]
                    f.write(f"{int(val)}\n")
    
    print(f"Saved output (text) to {filename}")


if __name__ == "__main__":
    # Configuration matching the deconv module
    # For deconv: output_size = (input_size - 1) * stride + kernel - 2*padding + output_padding
    # Example: (8 - 1) * 2 + 3 - 2*1 + 0 = 14 + 3 - 2 = 15
    
    input_height, input_width = 8, 8  # Input size for deconv
    in_channels = 16
    out_channels = 16
    
    P_ICH, P_OCH = 4, 4  # Parallel channels
    kernel_size = 3
    stride = 2
    padding = 1
    output_padding = 0
    
    A_BIT, W_BIT, B_BIT = 8, 8, 32
    
    # Calculate output dimensions
    out_height = (input_height - 1) * stride + kernel_size - 2 * padding + output_padding
    out_width = (input_width - 1) * stride + kernel_size - 2 * padding + output_padding

    # Generate test data using arange (0-127 range)
    # Input data
    input_size = input_height * input_width * in_channels
    input_data = (np.arange(input_size) % 128).astype(np.uint8)  # 0-127, uint8
    input_data = input_data.reshape(input_height, input_width, in_channels)
    
    # Weight data
    weight_size = out_channels * in_channels * kernel_size * kernel_size
    weight_data = (np.arange(weight_size) % 128).astype(np.uint8)  # 0-127, uint8
    weight = weight_data.reshape(out_channels, in_channels, kernel_size, kernel_size)
    
    bias = np.zeros(out_channels, dtype=np.int32)

    print("="*60)
    print("Deconvolution Configuration:")
    print(f"  Input: {input_height}x{input_width}x{in_channels}")
    print(f"  Weight: {out_channels}x{in_channels}x{kernel_size}x{kernel_size}")
    print(f"  Output: {out_height}x{out_width}x{out_channels}")
    print(f"  P_ICH={P_ICH}, P_OCH={P_OCH}")
    print(f"  Kernel={kernel_size}, Stride={stride}, Padding={padding}, Output_Padding={output_padding}")
    print("="*60)

    # Run deconvolution
    output = deconvolution(input_data, weight, bias, stride, padding, output_padding)

    print(f"\nOutput shape: {output.shape}")
    print(f"Output min: {np.min(output):.2f}, max: {np.max(output):.2f}, mean: {np.mean(output):.2f}")

    # Reshape weight for hardware
    print("\n" + "="*60)
    print("Reshaping weight for hardware...")
    weight_hw = reshape_weight_for_hardware(weight, P_OCH, P_ICH)
    print(f"Weight hardware shape: {weight_hw.shape}")
    print(f"  (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH) = {weight_hw.shape}")

    # Save to binary files
    print("\n" + "="*60)
    print("Saving data to binary and text files...")
    
    # Save parameters to SystemVerilog header file
    with open('deconv_config.svh', 'w') as f:
        f.write("// Auto-generated configuration file for deconv testbench\n")
        f.write("// Generated by gen_deconv.py\n\n")
        f.write(f"localparam int unsigned P_ICH = {P_ICH};\n")
        f.write(f"localparam int unsigned P_OCH = {P_OCH};\n")
        f.write(f"localparam int unsigned N_ICH = {in_channels};\n")
        f.write(f"localparam int unsigned N_OCH = {out_channels};\n")
        f.write(f"localparam int unsigned K = {kernel_size};\n")
        f.write(f"localparam int unsigned A_BIT = {A_BIT};\n")
        f.write(f"localparam int unsigned W_BIT = {W_BIT};\n")
        f.write(f"localparam int unsigned B_BIT = {B_BIT};\n")
        f.write(f"localparam int unsigned STRIDE = {stride};\n")
        f.write(f"localparam int unsigned PADDING = {padding};\n")
        f.write(f"localparam int unsigned OUTPUT_PADDING = {output_padding};\n")
        f.write(f"localparam int unsigned IN_H = {input_height};\n")
        f.write(f"localparam int unsigned IN_W = {input_width};\n")
        f.write(f"localparam int unsigned OUT_H = {out_height};\n")
        f.write(f"localparam int unsigned OUT_W = {out_width};\n")
        f.write(f"\n// Derived parameters\n")
        f.write(f"localparam int unsigned FOLD_I = N_ICH / P_ICH;\n")
        f.write(f"localparam int unsigned FOLD_O = N_OCH / P_OCH;\n")
        f.write(f"localparam int unsigned KK = K * K;\n")
        f.write(f"localparam int unsigned WEIGHT_DEPTH = FOLD_O * FOLD_I * KK;\n")
        f.write(f"localparam int unsigned IN_HW = IN_H * IN_W;\n")
        f.write(f"localparam int unsigned OUT_HW = OUT_H * OUT_W;\n")
    print("Saved configuration to deconv_config.svh")
    
    # Save input
    save_input_to_binary(input_data, 'deconv_input.bin', P_ICH, A_BIT)
    save_input_to_text(input_data, 'deconv_input.txt')
    
    # Save weight
    save_weight_to_binary(weight_hw, 'deconv_weight.bin', P_OCH, P_ICH, W_BIT)
    save_weight_to_text(weight_hw, 'deconv_weight.txt')
    
    # Save output
    save_output_to_binary(output, 'deconv_output.bin', P_OCH, B_BIT)
    save_output_to_text(output, 'deconv_output.txt')
